import time
import logging
import re
from airtest.core.api import stop_app, connect_device, snapshot
from airtest.core.android import Android
from utils import retry
from airtest.core.android.adb import ADB
from config import CLEANUP_DIR, BASELINE_PATH
from utils import compare_images  # 依赖工具函数
from config import INTERRUPTED

logger = logging.getLogger(__name__)

# 已知设备映射表（型号: URI）
DEVICE_MAPPING = {
    "HUAWEI ABR-AL00": "Android://127.0.0.1:5037/YSE0222325008504",
    "VIVO V2020CA": "android://127.0.0.1:5037/d8778bea",
    "XIAOMI M2102K1AC":"android://127.0.0.1:5037/673264a4",
    "Oppo find X3":"android://127.0.0.1:5037/de4f7957",
    "XIAOMI Mi 10":"android://127.0.0.1:5037/bc06c4a"
}


def get_device_model(serialno: str) -> str:
    """通过ADB获取设备型号"""
    try:
        adb = ADB(serialno)
        model = adb.shell("getprop ro.product.model").strip()
        brand = adb.shell("getprop ro.product.brand").strip().upper()
        return f"{brand} {model}"
    except Exception as e:
        logger.error(f"获取设备型号失败: {str(e)}")
        return "UNKNOWN"


def auto_select_device() -> str:
    """自动选择设备URI"""
    try:
        # 获取所有连接设备
        adb = ADB()
        adb_devices = adb.devices()
        if not adb_devices:
            raise RuntimeError("未检测到连接的Android设备")

        # 选择第一个设备
        serialno = adb_devices[0][0]

        # 获取设备型号
        model = get_device_model(serialno)
        logger.info(f"检测到设备: {model} (序列号: {serialno})")

        # 在映射表中查找
        for known_model, uri in DEVICE_MAPPING.items():
            if known_model.upper() in model.upper():
                logger.info(f"匹配到已知设备: {known_model}")
                return uri

        # 未匹配到已知设备，使用默认连接方式
        logger.warning(f"未识别的设备型号: {model}, 使用默认连接方式")
        return f"android://127.0.0.1:5037/{serialno}"

    except Exception as e:
        logger.error(f"自动选择设备失败: {str(e)}")
        raise


def connect_device_android() -> Android:
    """连接Android设备（自动选择）"""
    try:
        # 自动获取设备URI
        uri = auto_select_device()
        dev: Android = connect_device(uri)
        logger.info(f"✅ 设备连接成功: {dev.serialno}")
        return dev
    except Exception as e:
        logger.error(f"❌ 设备连接失败: {str(e)}")
        raise


def perform_device_cleanup(dev: Android) -> bool:
    """执行环境清理"""
    try:
        timestamp = time.strftime("%Y%m%d%H%M%S")

        # 清理主应用
        logger.info("🚨 强制清理主应用...")
        stop_app("cn.migu.video.clip.walle")
        time.sleep(2)

        # 保存清理后截图
        cleaned_shot = CLEANUP_DIR / f"cleaned_home_{timestamp}.png"
        dev.snapshot(str(cleaned_shot))
        logger.info(f"清理后主页截图：{cleaned_shot}")

        # 验证基准图匹配
        if compare_images(BASELINE_PATH, cleaned_shot):
            logger.success("环境恢复成功（主页一致）")
            return True

        # 备选清理（相册）
        logger.warning("主应用清理未完全恢复，尝试清理相册...")
        stop_app("com.huawei.photos")
        time.sleep(2)
        after_album_clean = CLEANUP_DIR / f"after_album_clean_{timestamp}.png"
        dev.snapshot(str(after_album_clean))
        if compare_images(BASELINE_PATH, after_album_clean):
            logger.success("通过清理相册恢复主页")
            return True

        logger.error("双重清理后仍未恢复主页")
        return False
    except Exception as e:
        logger.error(f"清理过程异常: {str(e)}")
        return False

class DeviceManager:
    @retry(times=3, delay=2)
    def shell(self, command):
        """执行ADB shell命令并返回输出"""
        return self.dev.shell(command)